import open3d as o3d
import numpy as np


def extract_features(
    pcd,
    voxel_size=0.02,
    normal_radius=0.1,
    normal_max_nn=30,
    curvature_knn=10,
    angle_threshold=0.5,
    curvature_threshold=0.2,
):
    """
    对输入点云进行去噪、降采样、法线估计、曲率计算和特征点提取，返回特征点云
    """
    print("[1/5]去噪...")
    cl, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
    pcd = pcd.select_by_index(ind)
    print(f"去噪后点数: {np.asarray(pcd.points).shape[0]}")

    print("[2/5]降采样...")
    pcd_down = pcd.voxel_down_sample(voxel_size)
    print(f"降采样后点数: {np.asarray(pcd_down.points).shape[0]}")

    print("[3/5]估计法线...")
    pcd_down.estimate_normals(
        search_param=o3d.geometry.KDTreeSearchParamHybrid(
            radius=normal_radius, max_nn=normal_max_nn
        )
    )
    points = np.asarray(pcd_down.points)
    normals = np.asarray(pcd_down.normals)

    print("[4/5]计算曲率...")
    curvatures = []
    kdtree = o3d.geometry.KDTreeFlann(pcd_down)
    for i in range(len(points)):
        [k, idx, _] = kdtree.search_knn_vector_3d(pcd_down.points[i], curvature_knn + 1)
        neighbor_points = points[idx[1:], :]
        neighbor_centered = neighbor_points - points[i]
        cov = np.cov(neighbor_centered, rowvar=False)
        eigenvalues, _ = np.linalg.eig(cov)
        if eigenvalues.sum() > 0:
            curvature = eigenvalues.min() / eigenvalues.sum()
        else:
            curvature = 0
        curvatures.append(curvature)
    curvatures = np.array(curvatures)

    print("[5/5]提取特征点...")
    feature_indices = []
    for i in range(len(points)):
        [k, idx, _] = kdtree.search_knn_vector_3d(pcd_down.points[i], 30)
        neighbor_normals = normals[idx[1:], :]
        current_normal = normals[i]
        dot_products = np.abs(np.dot(neighbor_normals, current_normal))
        angles = np.arccos(np.clip(dot_products, 0, 1))
        max_angle = np.max(angles) if len(angles) > 0 else 0
        if max_angle > angle_threshold and curvatures[i] > curvature_threshold:
            feature_indices.append(i)
    print(f"特征点数: {len(feature_indices)}")

    feature_pcd = pcd_down.select_by_index(feature_indices)
    return feature_pcd


if __name__ == "__main__":
    # ==== 参数配置 ====
    input_path = "../maps/stair1_clip.pcd"  # 输入点云文件路径
    output_path = "/maps/stair1_clip_feature.pcd"  # 输出特征点云文件路径
    voxel_size = 0.02  # 降采样体素大小
    normal_radius = 0.08  # 法线估计半径
    normal_max_nn = 30  # 法线估计最大邻域点数
    curvature_knn = 10  # 曲率计算邻域点数
    angle_threshold = 0.4  # 法线夹角阈值（弧度，约28.6度）
    curvature_threshold = 0.3  # 曲率阈值（需根据数据调整）
    # ===================

    print(f"读取点云文件: {input_path}")
    pcd = o3d.io.read_point_cloud(input_path)
    print(f"原始点数: {np.asarray(pcd.points).shape[0]}")

    # 特征提取
    feature_pcd = extract_features(
        pcd,
        voxel_size=voxel_size,
        normal_radius=normal_radius,
        normal_max_nn=normal_max_nn,
        curvature_knn=curvature_knn,
        angle_threshold=angle_threshold,
        curvature_threshold=curvature_threshold,
    )

    # 保存与可视化
    o3d.io.write_point_cloud(output_path, feature_pcd)
    print(f"特征点云已保存到: {output_path}")

    o3d.visualization.draw_geometries([feature_pcd])
